import numpy as np
import pymysql
import concurrent.futures
from datetime import datetime

class DatabaseManager:
    def __init__(self, database):
        if database == "mydatabase":
            self.db = pymysql.connect(host='127.0.0.1', user='root', passwd='wla', db=database, charset='utf8')
        if database == "manhattan_db":
            self.db = pymysql.connect(host='47.94.239.94', port=3306, user='manhattan', db=database, password='123456')
        self.cursor = self.db.cursor()

    def store_row_data(self, row_data, cursor):
        try:
            state_dt = (datetime.strptime(str(row_data['trade_date']), "%Y%m%d")).strftime('%Y-%m-%d')
            sql_insert = "INSERT INTO stock_daily(state_dt, stock_code, open, high, low, close, pre_close, change_val, pct_chg, vol, amount) VALUES" \
                         "('%s', '%s', '%.2f', '%.2f','%.2f','%.2f', '%.2f', '%.2f', '%.2f', '%.2f', '%.2f')" % \
                         (state_dt, str(row_data['ts_code']), float(row_data['open']), float(row_data['high']), float(row_data['low']),
                          float(row_data['close']), float(row_data['pre_close']), float(row_data['change']),
                          float(row_data['pct_chg']), float(row_data['vol']), float(row_data['amount']))
            cursor.execute(sql_insert)
        except Exception as e:
            print(f"Error occurred: {str(e)}")

    def store_data(self, processed_data: dict):
        data = processed_data['data']
        if processed_data['type'] == 'daily':
            table_exists = False
            try:
                self.cursor.execute("SELECT 1 FROM stock_daily LIMIT 1")
                table_exists = True
            except pymysql.Error as e:
                if e.args[0] == 1146:  # 表不存在的错误码
                    table_exists = False

            if not table_exists:
                create_table_query = """
                CREATE TABLE stock_daily (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    state_dt VARCHAR(50),
                    stock_code VARCHAR(50),
                    open FLOAT,
                    high FLOAT,
                    low FLOAT,
                    close FLOAT,
                    pre_close FLOAT,
                    change_val FLOAT,
                    pct_chg FLOAT,
                    vol FLOAT,
                    amount FLOAT
                );
                """
                self.cursor.execute(create_table_query)

            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
                # 提交存储数据的任务给线程池
                futures = [executor.submit(self.store_row_data, row_data, self.cursor)
                           for index, row_data in data.iterrows()]
                # 等待所有任务完成
                concurrent.futures.wait(futures)
            self.db.commit()
            self.cursor.close()
            self.db.close()

    def fetch_data(self, table, turnover_table, start_date, end_date) -> dict:
        start_date = datetime.strptime(start_date, "%Y-%m-%d").strftime("%Y-%m-%d")
        end_date = datetime.strptime(end_date, "%Y-%m-%d").strftime("%Y-%m-%d")

        query = "SELECT * FROM stock_daily WHERE STR_TO_DATE(state_dt, '%Y-%m-%d') BETWEEN '{}' AND '{}'".format(start_date, end_date)
        self.cursor.execute(query)
        results = self.cursor.fetchall()

        query_turnover = f"SELECT * FROM {turnover_table}"
        self.cursor.execute(query_turnover)
        turnover_results = self.cursor.fetchall()
        turnover = {}
        for row in turnover_results:
            stock_code = row[2]
            if stock_code not in turnover:
                turnover[stock_code] = {}
            turnover[stock_code][row[1]] = row[3]

        # 构造字典来存储数据
        stock_data_dict = {}
        for row in results:
            stock_code = row[2]
            if stock_code not in stock_data_dict:
                stock_data_dict[stock_code] = {}
            stock_data_dict[stock_code][row[1]] = {'open': row[3], 'high': row[4], 'low': row[5], 'close': row[6],
                                                   'pre_close': row[7], 'change': row[8], 'pct_chg': row[9],
                                                   'vol': row[10], 'amount': row[11]}
            if stock_code in turnover and row[1] in turnover[stock_code]:
                stock_data_dict[stock_code][row[1]]['turnover_rate'] = turnover[stock_code][row[1]]
        self.cursor.close()
        self.db.close()
        return stock_data_dict

    # 存储单条
    def store_single_turnover(self, row_data, cursor):
        state_dt = (datetime.strptime(str(row_data['trade_date']), "%Y%m%d")).strftime('%Y-%m-%d')
        sql_insert = "INSERT INTO turnover_daily(state_dt, stock_code, turnover_rate) VALUES" \
                     "('%s', '%s', '%.2f')" % \
                     (state_dt, str(row_data['ts_code']), float(row_data['turnover_rate']))
        cursor.execute(sql_insert)

    def store_turnover(self, processed_data: dict):
        data = processed_data['data']
        table_exists = False
        try:
            self.cursor.execute("SELECT 1 FROM turnover_daily LIMIT 1")
            table_exists = True
        except pymysql.Error as e:
            if e.args[0] == 1146:  # 表不存在的错误码
                table_exists = False
        if not table_exists:
            create_table_query = """
            CREATE TABLE turnover_daily (
                state_dt VARCHAR(50),
                stock_code VARCHAR(50),
                turnover_rate FLOAT,
                PRIMARY KEY (state_dt, stock_code)
            );
            """
            self.cursor.execute(create_table_query)
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            futures = [executor.submit(self.store_single_turnover, row_data, self.cursor)
                       for index, row_data in data.iterrows()]
            concurrent.futures.wait(futures)
        self.db.commit()
        self.cursor.close()
        self.db.close()

    def fetch_manhattan(self, start_date, end_date, detect_date):
        detect_date_query = "SELECT ts_code FROM stock_daily_info WHERE trade_date = '{}' AND pct_chg > 9.9".format(
            detect_date)
        self.cursor.execute(detect_date_query)
        ts_codes_result = self.cursor.fetchall()
        try:
            ts_code_list = [row[0] for row in ts_codes_result]
            ts_code_string = ', '.join(['%s'] * len(ts_code_list))
            start_end_date_query = "SELECT * FROM stock_daily_info WHERE trade_date BETWEEN '{}' AND '{}' AND ts_code IN ({})".format(
                start_date, end_date, ts_code_string)
            self.cursor.execute(start_end_date_query, ts_code_list)
            results = self.cursor.fetchall()
            # 构造字典来存储数据
            stock_data_dict = {}
            for row in results:
                stock_code = row[1]
                stock_name = row[2]
                trade_date = row[3]
                if stock_code not in stock_data_dict:
                    stock_data_dict[stock_code] = {}
                stock_data_dict[stock_code][trade_date] = {'name': stock_name, 'open': row[4], 'high': row[5], 'low': row[6], 'close': row[7],
                                                       'pre_close': row[8], 'change': row[9], 'pct_chg': row[10],
                                                       'vol': row[11], 'amount': row[12], 'vol_ratio': row[13], 'turn_over': row[14]}
            self.cursor.close()
            self.db.close()
            return stock_data_dict
        except Exception as e:
            print(e)
